import random
from prob import *


def create_train_data_list_random(num, my_filter, bound_min, bound_max): # 1024
    cnt = 0
    data_list = []
    while num - cnt != 0:
        curr_x = []
        for i_dim in range(DIM):
            xi = random.uniform(bound_min[i_dim], bound_max[i_dim])
            curr_x.append(xi)
        if my_filter(curr_x):
            cnt += 1
            data_list.append(curr_x)
    return data_list


def create_train_data_list(num, my_filter, bound_min, bound_max):
    data_list = []
    cnt = 0
    step_size = [0.] * DIM
    for i_dim in range(DIM):
        step_size[i_dim] = (bound_max[i_dim] - bound_min[i_dim]) / num[i_dim]

    # 网格点
    # actually, points have 129*129
    for a0 in range(num[0] + 1):
        x0 = bound_min[0] + a0 * step_size[0]
        for a1 in range(num[1] + 1):
            x1 = bound_min[1] + a1 * step_size[1]
            if DIM == 2:
                curr_x = [x0, x1]
                if my_filter(curr_x):
                    cnt += 1
                    data_list.append(curr_x)
            else:
                for a2 in range(num[2] + 1):
                    x2 = bound_min[2] + a2 * step_size[2]
                    if DIM == 3:
                        curr_x = [x0, x1, x2]
                        if my_filter(curr_x):
                            cnt += 1
                            data_list.append(curr_x)
                    else:
                        for a3 in range(num[3] + 1):
                            x3 = bound_min[3] + a3 * step_size[3]
                            if DIM == 4:
                                curr_x = [x0, x1, x2, x3]
                                if my_filter(curr_x):
                                    cnt += 1
                                    data_list.append(curr_x)
                            else:
                                for a4 in range(num[4] + 1):
                                    x4 = bound_min[4] + a4 * step_size[4]
                                    if DIM == 5:
                                        curr_x = [x0, x1, x2, x3, x4]
                                        if my_filter(curr_x):
                                            cnt += 1
                                            data_list.append(curr_x)
                                    else:
                                        for a5 in range(num[5] + 1):
                                            x5 = bound_min[5] + a5 * step_size[5]
                                            if DIM == 6:
                                                curr_x = [x0, x1, x2, x3, x4, x5]
                                                if my_filter(curr_x):
                                                    cnt += 1
                                                    data_list.append(curr_x)
                                            else:
                                                for a6 in range(num[6] + 1):
                                                    x6 = bound_min[6] + a6 * step_size[6]
                                                    if DIM == 7:
                                                        curr_x = [x0, x1, x2, x3, x4, x5, x6]
                                                        if my_filter(curr_x):
                                                            cnt += 1
                                                            data_list.append(curr_x)
                                                    else:
                                                        for a7 in range(num[7] + 1):
                                                            x7 = bound_min[7] + a7 * step_size[7]
                                                            if DIM == 8:
                                                                curr_x = [x0, x1, x2, x3, x4, x5, x6, x7]
                                                                if my_filter(curr_x):
                                                                    cnt += 1
                                                                    data_list.append(curr_x)
                                                            else:
                                                                for a8 in range(num[8] + 1):
                                                                    x8 = bound_min[8] + a8 * step_size[8]
                                                                    if DIM == 9:
                                                                        curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8]
                                                                        if my_filter(curr_x):
                                                                            cnt += 1
                                                                            data_list.append(curr_x)
                                                                    else:
                                                                        for a9 in range(num[9] + 1):
                                                                            x9 = bound_min[9] + a9 * step_size[9]
                                                                            if DIM == 10:
                                                                                curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9]
                                                                                if my_filter(curr_x):
                                                                                    cnt += 1
                                                                                    data_list.append(curr_x)
                                                                            else:
                                                                                for a10 in range(num[10] + 1):
                                                                                    x10 = bound_min[10] + a10 * step_size[10]
                                                                                    if DIM == 11:
                                                                                        curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10]
                                                                                        if my_filter(curr_x):
                                                                                            cnt += 1
                                                                                            data_list.append(curr_x)
                                                                                    else:
                                                                                        for a11 in range(num[11] + 1):
                                                                                            x11 = bound_min[11] + a11 * step_size[11]
                                                                                            if DIM == 12:
                                                                                                curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11]
                                                                                                if my_filter(curr_x):
                                                                                                    cnt += 1
                                                                                                    data_list.append(curr_x)
                                                                                            else:
                                                                                                for a12 in range(num[12] + 1):
                                                                                                    x12 = bound_min[12] + a12 * time_step[12]
                                                                                                    if DIM == 13:
                                                                                                        curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12]
                                                                                                        if my_filter(curr_x):
                                                                                                            cnt += 1
                                                                                                            data_list.append(curr_x)
                                                                                                    else:
                                                                                                        for a13 in range(num[13] + 1):
                                                                                                            x13 = bound_min[13] + a13 * time_step[13]
                                                                                                            if DIM == 14:
                                                                                                                curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13]
                                                                                                                if my_filter(curr_x):
                                                                                                                    cnt += 1
                                                                                                                    data_list.append(curr_x)
                                                                                                            else:
                                                                                                                for a14 in range(num[14] + 1):
                                                                                                                    x14 = bound_min[14] + a14 * time_step[14]
                                                                                                                    if DIM == 15:
                                                                                                                        curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14]
                                                                                                                        if my_filter(curr_x):
                                                                                                                            cnt += 1
                                                                                                                            data_list.append(curr_x)
                                                                                                                    else:
                                                                                                                        for a15 in range(num[15] + 1):
                                                                                                                            x15 = bound_min[15] + a15 * time_step[15]
                                                                                                                            if DIM == 16:
                                                                                                                                curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15]
                                                                                                                                if my_filter(curr_x):
                                                                                                                                    cnt += 1
                                                                                                                                    data_list.append(curr_x)
                                                                                                                            else:
                                                                                                                                for a16 in range(num[16] + 1):
                                                                                                                                    x16 = bound_min[16] + a16 * time_step[16]
                                                                                                                                    if DIM == 17:
                                                                                                                                        curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15, x16]
                                                                                                                                        if my_filter(curr_x):
                                                                                                                                            cnt += 1
                                                                                                                                            data_list.append(curr_x)
                                                                                                                                    else:
                                                                                                                                        for a17 in range(num[17] + 1):
                                                                                                                                            x17 = bound_min[17] + a17 * time_step[17]
                                                                                                                                            if DIM == 18:
                                                                                                                                                curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15, x16, x17]
                                                                                                                                                if my_filter(curr_x):
                                                                                                                                                    cnt += 1
                                                                                                                                                    data_list.append(curr_x)
                                                                                                                                            else:
                                                                                                                                                for a18 in range(num[18] + 1):
                                                                                                                                                    x18 = bound_min[18] + a18 * time_step[18]
                                                                                                                                                    if DIM == 19:
                                                                                                                                                        curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15, x16, x17, x18]
                                                                                                                                                        if my_filter(curr_x):
                                                                                                                                                            cnt += 1
                                                                                                                                                            data_list.append(curr_x)
                                                                                                                                                    else:
                                                                                                                                                        print("can only compute dim from 2 to 19!")


    # 随机点
    # random_points_num = 10000
    # while random_points_num > 0:
    #     curr_x = []
    #     for i_dim in range(DIM):
    #         xi = random.uniform(bound_min[i_dim], bound_max[i_dim])
    #         curr_x.append(xi)
    #     if my_filter(curr_x):
    #         cnt += 1
    #         random_points_num -= 1
    #         data_list.append(curr_x)

    # 不足一小批次的补齐
    while cnt % BATCH_SIZE != 0:
        curr_x = []
        for i_dim in range(DIM):
            xi = random.uniform(bound_min[i_dim], bound_max[i_dim])
            curr_x.append(xi)
        if my_filter(curr_x):
            cnt += 1
            data_list.append(curr_x)
    return data_list


def generate_data():
    if DIM <= 10:
        batches_init = create_train_data_list(piece_num[0], is_init, init_min, init_max)
        batches_unsafe = create_train_data_list(piece_num[1], is_unsafe, unsafe_min, unsafe_max)
        batches_domain = create_train_data_list(piece_num[2], is_domain, domain_min, domain_max)
    else:
        batches_init = create_train_data_list_random(piece_num[0], is_init, init_min, init_max)
        batches_unsafe = create_train_data_list_random(piece_num[1], is_unsafe, unsafe_min, unsafe_max)
        batches_domain = create_train_data_list_random(piece_num[2], is_domain, domain_min, domain_max)
    return batches_init, batches_unsafe, batches_domain


def generate_data_eg2():
    batches_init0 = create_train_data_list([9, 9], is_init, sub_init[0][0], sub_init[0][1])
    batches_init1 = create_train_data_list([9, 9], is_init, sub_init[1][0], sub_init[1][1])
    batches_init2 = create_train_data_list([9, 9], is_init, sub_init[2][0], sub_init[2][1])
    batches_unsafe0 = create_train_data_list([9, 9], is_unsafe, sub_unsafe[0][0], sub_unsafe[0][1])
    batches_unsafe1 = create_train_data_list([9, 9], is_unsafe, sub_unsafe[1][0], sub_unsafe[1][1])
    batches_unsafe2 = create_train_data_list([9, 9], is_unsafe, sub_unsafe[2][0], sub_unsafe[2][1])
    batches_domain = create_train_data_list(piece_num[2], is_domain, domain_min, domain_max)

    batches_init = batches_init0 + batches_init1 + batches_init2
    batches_unsafe = batches_unsafe0 + batches_unsafe1 + batches_unsafe2
    return batches_init, batches_unsafe, batches_domain

def generate_data_eg6():
    batches_init = create_train_data_list(piece_num[0], is_init, init_min, init_max)
    batches_unsafe0 = create_train_data_list([9, 9], is_unsafe, sub_unsafe[0][0], sub_unsafe[0][1])
    batches_unsafe1 = create_train_data_list([9, 9], is_unsafe, sub_unsafe[1][0], sub_unsafe[1][1])
    batches_unsafe2 = create_train_data_list([9, 9], is_unsafe, sub_unsafe[2][0], sub_unsafe[2][1])
    batches_unsafe3 = create_train_data_list([9, 9], is_unsafe, sub_unsafe[3][0], sub_unsafe[3][1])
    batches_domain = create_train_data_list(piece_num[2], is_domain, domain_min, domain_max)
    batches_unsafe = batches_unsafe0 + batches_unsafe1 + batches_unsafe2 + batches_unsafe3
    return batches_init, batches_unsafe, batches_domain

def add_cex(num, batches_data, cexes, my_filter):
    cnt = 0
    for cex in cexes:
        ori_min = []
        for i_dim in range(DIM):
            ori_min.append(cex[i_dim] - cex_step * int(num[i_dim]/2))
        ori_max = []
        for i_dim in range(DIM):
            ori_max.append(cex[i_dim] + cex_step * int(num[i_dim]/2))

        for a0 in range(num[0]):
            x0 = ori_min[0] + cex_step * a0
            for a1 in range(num[1]):
                x1 = ori_min[1] + cex_step * a1
                if DIM == 2:
                    curr_x = [x0, x1]
                    if my_filter(curr_x):
                        cnt += 1
                        batches_data.append(curr_x)
                else:
                    for a2 in range(num[2]):
                        x2 = ori_min[2] + cex_step * a2
                        if DIM == 3:
                            curr_x = [x0, x1, x2]
                            if my_filter(curr_x):
                                cnt += 1
                                batches_data.append(curr_x)
                        else:
                            for a3 in range(num[3]):
                                x3 = ori_min[3] + cex_step * a3
                                if DIM == 4:
                                    curr_x = [x0, x1, x2, x3]
                                    if my_filter(curr_x):
                                        cnt += 1
                                        batches_data.append(curr_x)
                                else:
                                    for a4 in range(num[4]):
                                        x4 = ori_min[4] + cex_step * a4
                                        if DIM == 5:
                                            curr_x = [x0, x1, x2, x3, x4]
                                            if my_filter(curr_x):
                                                cnt += 1
                                                batches_data.append(curr_x)
                                        else:
                                            for a5 in range(num[5]):
                                                x5 = ori_min[5] + cex_step * a5
                                                if DIM == 6:
                                                    curr_x = [x0, x1, x2, x3, x4, x5]
                                                    if my_filter(curr_x):
                                                        cnt += 1
                                                        batches_data.append(curr_x)
                                                else:
                                                    for a6 in range(num[6]):
                                                        x6 = ori_min[6] + cex_step * a6
                                                        if DIM == 7:
                                                            curr_x = [x0, x1, x2, x3, x4, x5, x6]
                                                            if my_filter(curr_x):
                                                                cnt += 1
                                                                batches_data.append(curr_x)
                                                        else:
                                                            for a7 in range(num[7]):
                                                                x7 = ori_min[7] + cex_step * a7
                                                                if DIM == 8:
                                                                    curr_x = [x0, x1, x2, x3, x4, x5, x6, x7]
                                                                    if my_filter(curr_x):
                                                                        cnt += 1
                                                                        batches_data.append(curr_x)
                                                                else:
                                                                    for a8 in range(num[8]):
                                                                        x8 = ori_min[8] + cex_step * a8
                                                                        if DIM == 9:
                                                                            curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8]
                                                                            if my_filter(curr_x):
                                                                                cnt += 1
                                                                                batches_data.append(curr_x)
                                                                        else:
                                                                            for a9 in range(num[9]):
                                                                                x9 = ori_min[9] + cex_step * a9
                                                                                if DIM == 10:
                                                                                    curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9]
                                                                                    if my_filter(curr_x):
                                                                                        cnt += 1
                                                                                        batches_data.append(curr_x)
                                                                                else:
                                                                                    for a10 in range(num[10]):
                                                                                        x10 = ori_min[10] + cex_step * a10
                                                                                        if DIM == 11:
                                                                                            curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10]
                                                                                            if my_filter(curr_x):
                                                                                                cnt += 1
                                                                                                batches_data.append(curr_x)
                                                                                        else:
                                                                                            for a11 in range(num[11]):
                                                                                                x11 = ori_min[11] + cex_step * a11
                                                                                                if DIM == 12:
                                                                                                    curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11]
                                                                                                    if my_filter(curr_x):
                                                                                                        cnt += 1
                                                                                                        batches_data.append(curr_x)
                                                                                                else:
                                                                                                    for a12 in range(num[12]):
                                                                                                        x12 = ori_min[12] + cex_step * a12
                                                                                                        if DIM == 13:
                                                                                                            curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12]
                                                                                                            if my_filter(curr_x):
                                                                                                                cnt += 1
                                                                                                                batches_data.append(curr_x)
                                                                                                        else:
                                                                                                            for a13 in range(num[13]):
                                                                                                                x13 = ori_min[13] + cex_step * a13
                                                                                                                if DIM == 14:
                                                                                                                    curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13]
                                                                                                                    if my_filter(curr_x):
                                                                                                                        cnt += 1
                                                                                                                        batches_data.append(curr_x)
                                                                                                                else:
                                                                                                                    for a14 in range(num[14]):
                                                                                                                        x14 = ori_min[14] + cex_step * a14
                                                                                                                        if DIM == 15:
                                                                                                                            curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14]
                                                                                                                            if my_filter(curr_x):
                                                                                                                                cnt += 1
                                                                                                                                batches_data.append(curr_x)
                                                                                                                        else:
                                                                                                                            for a15 in range(num[15]):
                                                                                                                                x15 = ori_min[15] + cex_step * a15
                                                                                                                                if DIM == 16:
                                                                                                                                    curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15]
                                                                                                                                    if my_filter(curr_x):
                                                                                                                                        cnt += 1
                                                                                                                                        batches_data.append(curr_x)
                                                                                                                                else:
                                                                                                                                    for a16 in range(num[16]):
                                                                                                                                        x16 = ori_min[16] + cex_step * a16
                                                                                                                                        if DIM == 17:
                                                                                                                                            curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15, x16]
                                                                                                                                            if my_filter(curr_x):
                                                                                                                                                cnt += 1
                                                                                                                                                batches_data.append(curr_x)
                                                                                                                                        else:
                                                                                                                                            for a17 in range(num[17]):
                                                                                                                                                x17 = ori_min[17] + cex_step * a17
                                                                                                                                                if DIM == 18:
                                                                                                                                                    curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15, x16, x17]
                                                                                                                                                    if my_filter(curr_x):
                                                                                                                                                        cnt += 1
                                                                                                                                                        batches_data.append(curr_x)
                                                                                                                                                else:
                                                                                                                                                    for a18 in range(num[18]):
                                                                                                                                                        x18 = ori_min[18] + cex_step * a18
                                                                                                                                                        if DIM == 19:
                                                                                                                                                            curr_x = [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15, x16, x17, x18]
                                                                                                                                                            if my_filter(curr_x):
                                                                                                                                                                cnt += 1
                                                                                                                                                                batches_data.append(curr_x)
                                                                                                                                                        else:
                                                                                                                                                            print("can only compute dim from 2 to 19!")

        while cnt % BATCH_SIZE != 0:
            curr_x = []
            for i_dim in range(DIM):
                xi = random.uniform(ori_min[i_dim], ori_max[i_dim])
                curr_x.append(xi)
            if my_filter(curr_x):
                cnt += 1
                batches_data.append(curr_x)
    return batches_data


def add_cexes(batches_init, cex_init_curr, batches_unsafe, cex_unsafe_curr, batches_domain, cex_domain_curr):
    batches_init = add_cex(cex_num[0], batches_init, cex_init_curr, is_init)
    batches_unsafe = add_cex(cex_num[1], batches_unsafe, cex_unsafe_curr, is_unsafe)
    batches_domain = add_cex(cex_num[2], batches_domain, cex_domain_curr, is_domain)
    return batches_init, batches_unsafe, batches_domain